#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# 接收 /wrench_fake 话题，050 和 000 交替
# 如果接收到的力不超过期望力020，则控制机械臂往初始位置运动
# xyz方向力控制

import time
import numpy as np
import rospy
from moveit_commander import MoveGroupCommander
from geometry_msgs.msg import Pose, WrenchStamped

class ForceControl:
    def __init__(self):
        self.force_torque_filtered = np.zeros((3, 1))
        self.force_torque_reference = np.zeros((3, 1))
        self.force_torque_reference[1] = 2  # 参考力
        self.position_reference = np.zeros((3, 1))
        self.tcp_frame_cur_p = np.zeros((3, 1))
        self.tcp_frame_cur_M = np.eye(3)
        self.velocity = np.zeros((3, 1))
        self.acceleration = np.zeros((3, 1))
        self.K_position = np.eye(3) * 1.1
        self.B_position = np.eye(3) * 180
        self.M_position = np.eye(3) * 280
        self.dt = 1
        self.first_get_coord = True
        self.group = MoveGroupCommander("manipulator")

        # 订阅 /wrench_fake 话题
        self.wrench_sub = rospy.Subscriber('/wrench_fake', WrenchStamped, self.wrench_callback)
        
    def wrench_callback(self, msg):
        self.force_torque_filtered[0] = msg.wrench.force.x
        self.force_torque_filtered[1] = msg.wrench.force.y
        self.force_torque_filtered[2] = msg.wrench.force.z

    def update_current_pose(self):
        current_pose = self.group.get_current_pose().pose
        self.tcp_frame_cur_p[0, 0] = current_pose.position.x
        self.tcp_frame_cur_p[1, 0] = current_pose.position.y
        self.tcp_frame_cur_p[2, 0] = current_pose.position.z
        if self.first_get_coord:
            self.position_reference[0, 0] = current_pose.position.x
            self.position_reference[1, 0] = current_pose.position.y
            self.position_reference[2, 0] = current_pose.position.z
            self.first_get_coord = False

    def calculate_new_pose(self):
        force_difference = self.force_torque_filtered - self.force_torque_reference
        position_difference = self.tcp_frame_cur_p - self.position_reference
        Kp_position_term = np.dot(self.K_position, position_difference)
        B_velocity_term = np.dot(self.B_position, self.velocity)
        
        self.acceleration = np.dot(np.linalg.inv(self.M_position), 
                                   (force_difference - Kp_position_term - B_velocity_term))
        self.velocity += (self.acceleration * self.dt)
        delta_position = self.velocity * self.dt
        
        new_position = self.tcp_frame_cur_p - delta_position
        return new_position
    
    def force_below_threshold(self, threshold):
        return np.all(np.abs(self.force_torque_filtered) <= threshold)

def main():
    rospy.init_node('force_control_node')  # 初始化ROS节点
    group_name = "manipulator"
    move_group = MoveGroupCommander(group_name)
    force_control = ForceControl()
    
    # 设置初始位置
    initial_pose = Pose()
    initial_pose.position.x = 0
    initial_pose.position.y = 0.2516255937679144
    initial_pose.position.z = 0.2591765585666135
    initial_pose.orientation.w = 1.0
    
    move_group.set_pose_target(initial_pose)
    move_group.go(wait=True)
    move_group.stop()
    move_group.clear_pose_targets()
    
    try:
        while not rospy.is_shutdown():
            force_control.update_current_pose()
            
            # 检查接收到的力是否超过期望力
            if force_control.force_below_threshold(np.array([[0], [2], [0]])):
                # 如果接收到的力不超过期望力，控制机械臂运动到初始位置
                move_group.set_pose_target(initial_pose)
                move_group.go(wait=False)
            else:
                new_position = force_control.calculate_new_pose()
                
                pose_target = move_group.get_current_pose().pose
                pose_target.position.x = new_position[0, 0]
                pose_target.position.y = new_position[1, 0]
                pose_target.position.z = new_position[2, 0]
                
                move_group.set_pose_target(pose_target)
                move_group.go(wait=False)
            
            time.sleep(0.5)
    except KeyboardInterrupt:
        print("Motion stopped by user")
    finally:
        move_group.stop()
        move_group.clear_pose_targets()

if __name__ == "__main__":
    main()
